#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
高胜率股票预测算法 - 带止盈策略
买入 -> 持有 -> 达到盈利目标 -> 提示卖出
"""

import numpy as np
import pandas as pd
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass
from enum import Enum
import json

class Signal(Enum):
    """交易信号"""
    STRONG_BUY = "强烈买入"
    TAKE_PROFIT = "止盈卖出"
    HOLD = "持有观望"

class SignalType(Enum):
    """信号类型"""
    BUY = "BUY"
    SELL = "SELL"
    HOLD = "HOLD"

@dataclass
class TradingPoint:
    """交易点"""
    date: str
    price: float
    signal: Signal
    signal_type: SignalType
    confidence: float
    reasons: List[str]
    indicators: Dict[str, float]
    profit_rate: float = 0  # 盈利比例（卖出时）
    buy_price: float = 0  # 买入价格（卖出时）

class EnhancedTechnicalAnalysis:
    """技术分析"""
    
    @staticmethod
    def calculate_ma(prices: List[float], period: int) -> List[float]:
        """移动平均线"""
        if len(prices) < period:
            return []
        result = []
        for i in range(len(prices)):
            if i < period - 1:
                result.append(np.nan)
            else:
                result.append(np.mean(prices[i-period+1:i+1]))
        return result
    
    @staticmethod
    def calculate_ema(prices: List[float], period: int) -> List[float]:
        """指数移动平均"""
        if len(prices) < period:
            return []
        result = []
        multiplier = 2 / (period + 1)
        sma = np.mean(prices[:period])
        result.append(sma)
        for i in range(1, len(prices)):
            if i < period:
                result.append(np.nan)
            else:
                ema = (prices[i] - result[-1]) * multiplier + result[-1]
                result.append(ema)
        return result
    
    @staticmethod
    def calculate_macd(prices: List[float]) -> Dict[str, List[float]]:
        """MACD指标"""
        ema12 = EnhancedTechnicalAnalysis.calculate_ema(prices, 12)
        ema26 = EnhancedTechnicalAnalysis.calculate_ema(prices, 26)
        macd_line = [e12 - e26 if not np.isnan(e12) and not np.isnan(e26) else np.nan 
                     for e12, e26 in zip(ema12, ema26)]
        signal_line = EnhancedTechnicalAnalysis.calculate_ema(
            [x for x in macd_line if not np.isnan(x)], 9
        )
        signal_line = [np.nan] * (len(macd_line) - len(signal_line)) + signal_line
        histogram = [m - s if not np.isnan(m) and not np.isnan(s) else np.nan 
                     for m, s in zip(macd_line, signal_line)]
        return {'macd': macd_line, 'signal': signal_line, 'histogram': histogram}
    
    @staticmethod
    def calculate_rsi(prices: List[float], period: int = 14) -> List[float]:
        """RSI指标"""
        if len(prices) < period + 1:
            return []
        deltas = [prices[i] - prices[i-1] for i in range(1, len(prices))]
        gains = [d if d > 0 else 0 for d in deltas]
        losses = [-d if d < 0 else 0 for d in deltas]
        result = [np.nan]
        avg_gain = np.mean(gains[:period])
        avg_loss = np.mean(losses[:period])
        for i in range(period, len(deltas)):
            if avg_loss == 0:
                rsi = 100
            else:
                rs = avg_gain / avg_loss
                rsi = 100 - (100 / (1 + rs))
            result.append(rsi)
            avg_gain = (avg_gain * (period - 1) + gains[i]) / period
            avg_loss = (avg_loss * (period - 1) + losses[i]) / period
        return result
    
    @staticmethod
    def calculate_bollinger_bands(prices: List[float], period: int = 20) -> Dict[str, List[float]]:
        """布林带"""
        ma = EnhancedTechnicalAnalysis.calculate_ma(prices, period)
        upper, lower = [], []
        for i in range(len(prices)):
            if i < period - 1:
                upper.append(np.nan)
                lower.append(np.nan)
            else:
                std = np.std(prices[i-period+1:i+1])
                upper.append(ma[i] + 2 * std)
                lower.append(ma[i] - 2 * std)
        return {'upper': upper, 'middle': ma, 'lower': lower}

class ProfitTargetPredictor:
    """带止盈目标的预测器"""
    
    def __init__(self, 
                 min_buy_score: int = 8,
                 min_confidence: int = 70,
                 profit_target: float = 0.08,  # 止盈目标8%
                 trailing_stop: float = 0.03):  # 回撤止盈3%
        self.ta = EnhancedTechnicalAnalysis()
        self.min_buy_score = min_buy_score
        self.min_confidence = min_confidence
        self.profit_target = profit_target
        self.trailing_stop = trailing_stop
        self.current_position = None  # 当前持仓 {'buy_price': xx, 'buy_date': xx}
        self.max_profit_rate = 0  # 最大盈利率
    
    def analyze(self, stock_data: List[Dict]) -> List[TradingPoint]:
        """
        分析股票，生成买入和卖出信号
        策略：买入后达到盈利目标就提示卖出
        """
        if len(stock_data) < 60:
            return []
        
        dates = [d['date'] for d in stock_data]
        opens = [float(d['open']) for d in stock_data]
        closes = [float(d['close']) for d in stock_data]
        highs = [float(d['high']) for d in stock_data]
        lows = [float(d['low']) for d in stock_data]
        volumes = [float(d['volume']) for d in stock_data]
        
        # 计算技术指标
        ma5 = self.ta.calculate_ma(closes, 5)
        ma10 = self.ta.calculate_ma(closes, 10)
        ma20 = self.ta.calculate_ma(closes, 20)
        ma60 = self.ta.calculate_ma(closes, 60)
        macd = self.ta.calculate_macd(closes)
        rsi = self.ta.calculate_rsi(closes, 14)
        bollinger = self.ta.calculate_bollinger_bands(closes, 20)
        
        trading_points = []
        position = None  # 模拟持仓
        max_profit = 0
        
        for i in range(60, len(stock_data)):
            current_price = closes[i]
            
            # === 如果有持仓，检查是否止盈 ===
            if position:
                profit_rate = (current_price - position['buy_price']) / position['buy_price']
                max_profit = max(max_profit, profit_rate)
                
                should_sell = False
                sell_reasons = []
                
                # 1. 达到止盈目标（主要条件）
                if profit_rate >= self.profit_target:
                    should_sell = True
                    sell_reasons.append(f"达到止盈目标{profit_rate*100:.1f}%")
                
                # 2. 回撤止盈（从最高点回撤）
                elif max_profit > 0 and (max_profit - profit_rate) >= self.trailing_stop:
                    should_sell = True
                    sell_reasons.append(f"回撤止盈(最高{max_profit*100:.1f}%)")
                
                # 3. 技术指标恶化（辅助条件）
                if profit_rate > 0.03:  # 至少盈利3%时才考虑技术止盈
                    tech_sell_score = 0
                    
                    # RSI超买
                    if i < len(rsi) and not np.isnan(rsi[i]) and rsi[i] > 70:
                        tech_sell_score += 1
                        sell_reasons.append(f"RSI超买({rsi[i]:.1f})")
                    
                    # MACD死叉
                    if (i > 0 and not np.isnan(macd['macd'][i]) and 
                        macd['macd'][i-1] >= macd['signal'][i-1] and 
                        macd['macd'][i] < macd['signal'][i]):
                        tech_sell_score += 1
                        sell_reasons.append("MACD死叉")
                    
                    # MA死叉
                    if (i > 0 and not np.isnan(ma5[i]) and 
                        ma5[i-1] >= ma10[i-1] and ma5[i] < ma10[i]):
                        tech_sell_score += 1
                        sell_reasons.append("MA5死叉")
                    
                    if tech_sell_score >= 2:
                        should_sell = True
                
                # 生成卖出信号
                if should_sell:
                    confidence = min(80 + profit_rate * 100, 95)
                    
                    trading_point = TradingPoint(
                        date=dates[i],
                        price=current_price,
                        signal=Signal.TAKE_PROFIT,
                        signal_type=SignalType.SELL,
                        confidence=confidence,
                        reasons=sell_reasons[:3],
                        profit_rate=profit_rate,
                        buy_price=position['buy_price'],
                        indicators={
                            'ma5': ma5[i] if not np.isnan(ma5[i]) else 0,
                            'ma10': ma10[i] if not np.isnan(ma10[i]) else 0,
                            'rsi': rsi[i] if i < len(rsi) and not np.isnan(rsi[i]) else 0,
                        }
                    )
                    trading_points.append(trading_point)
                    position = None  # 清仓
                    max_profit = 0
                    continue
            
            # === 如果没有持仓，寻找买入信号 ===
            if not position:
                buy_score = 0
                buy_reasons = []
                
                # 1. MACD金叉（+3分）
                if (i > 0 and not np.isnan(macd['macd'][i]) and 
                    macd['macd'][i-1] <= macd['signal'][i-1] and 
                    macd['macd'][i] > macd['signal'][i] and
                    macd['macd'][i] < 0):
                    buy_score += 3
                    buy_reasons.append("MACD金叉（低位）")
                
                # 2. RSI超卖反弹（+3分）
                if i < len(rsi) and not np.isnan(rsi[i]):
                    if rsi[i] < 30 and i > 0 and rsi[i] > rsi[i-1]:
                        buy_score += 3
                        buy_reasons.append(f"RSI超卖反弹({rsi[i]:.1f})")
                
                # 3. 均线多头排列（+2分）
                if (not np.isnan(ma5[i]) and not np.isnan(ma10[i]) and 
                    not np.isnan(ma20[i]) and not np.isnan(ma60[i])):
                    if ma5[i] > ma10[i] > ma20[i]:
                        buy_score += 2
                        buy_reasons.append("多头排列")
                
                # 4. MA金叉（+2分）
                if (i > 0 and not np.isnan(ma5[i]) and 
                    ma5[i-1] <= ma10[i-1] and ma5[i] > ma10[i]):
                    buy_score += 2
                    buy_reasons.append("MA5金叉MA10")
                
                # 5. 布林带下轨反弹（+2分）
                if (not np.isnan(bollinger['lower'][i]) and 
                    closes[i] < bollinger['lower'][i] * 1.01 and 
                    i > 0 and closes[i] > closes[i-1]):
                    buy_score += 2
                    buy_reasons.append("下轨反弹")
                
                # 6. 价格低位（+1分）
                if not np.isnan(ma60[i]) and closes[i] < ma60[i] * 0.95:
                    buy_score += 1
                    buy_reasons.append("价格低位")
                
                # 7. 连续下跌后反弹（+1分）
                if i >= 3:
                    if all(closes[i-j] < closes[i-j-1] for j in range(3)) and closes[i] > closes[i-1]:
                        buy_score += 1
                        buy_reasons.append("连跌反弹")
                
                # 强烈买入信号
                if buy_score >= self.min_buy_score:
                    confidence = min(buy_score * 10, 95)
                    
                    if confidence >= self.min_confidence:
                        trading_point = TradingPoint(
                            date=dates[i],
                            price=current_price,
                            signal=Signal.STRONG_BUY,
                            signal_type=SignalType.BUY,
                            confidence=confidence,
                            reasons=buy_reasons[:4],
                            indicators={
                                'ma5': ma5[i] if not np.isnan(ma5[i]) else 0,
                                'ma10': ma10[i] if not np.isnan(ma10[i]) else 0,
                                'rsi': rsi[i] if i < len(rsi) and not np.isnan(rsi[i]) else 0,
                                'macd': macd['macd'][i] if not np.isnan(macd['macd'][i]) else 0,
                                'buy_score': buy_score
                            }
                        )
                        trading_points.append(trading_point)
                        position = {'buy_price': current_price, 'buy_date': dates[i]}
        
        return trading_points
    
    def get_current_recommendation(self, stock_data: List[Dict]) -> Dict:
        """获取当前建议（考虑持仓状态）"""
        if len(stock_data) < 60:
            return {
                'signal': 'HOLD',
                'signal_cn': '持有观望',
                'confidence': 0,
                'reasons': ['数据不足']
            }
        
        trading_points = self.analyze(stock_data)
        
        if not trading_points:
            return {
                'signal': 'HOLD',
                'signal_cn': '持有观望',
                'confidence': 50,
                'reasons': ['无强烈信号，建议观望'],
                'price': stock_data[-1]['close']
            }
        
        latest = trading_points[-1]
        return {
            'signal': latest.signal_type.value,
            'signal_cn': latest.signal.value,
            'confidence': latest.confidence,
            'reasons': latest.reasons,
            'price': latest.price,
            'date': latest.date,
            'profit_rate': latest.profit_rate * 100 if latest.profit_rate else 0,
            'buy_price': latest.buy_price if latest.buy_price else 0,
            'indicators': latest.indicators
        }
    
    def backtest(self, stock_data: List[Dict], initial_capital: float = 100000) -> Dict:
        """回测 - 买入后止盈卖出"""
        trading_points = self.analyze(stock_data)
        
        capital = initial_capital
        shares = 0
        trades = []
        win_trades = 0
        loss_trades = 0
        total_profit = 0
        
        for point in trading_points:
            if point.signal_type == SignalType.BUY and capital > 0:
                # 买入 - 使用80%资金
                invest = capital * 0.8
                shares = int(invest / point.price)
                if shares > 0:
                    cost = shares * point.price
                    capital -= cost
                    trades.append({
                        'date': point.date,
                        'action': 'BUY',
                        'price': point.price,
                        'shares': shares,
                        'confidence': point.confidence,
                        'reasons': point.reasons
                    })
            
            elif point.signal_type == SignalType.SELL and shares > 0:
                # 卖出
                revenue = shares * point.price
                capital += revenue
                
                # 计算本次交易盈亏
                if trades and trades[-1]['action'] == 'BUY':
                    buy_price = trades[-1]['price']
                    profit = (point.price - buy_price) / buy_price
                    total_profit += profit
                    
                    if profit > 0:
                        win_trades += 1
                    else:
                        loss_trades += 1
                
                trades.append({
                    'date': point.date,
                    'action': 'SELL',
                    'price': point.price,
                    'shares': shares,
                    'confidence': point.confidence,
                    'reasons': point.reasons,
                    'profit_rate': point.profit_rate * 100
                })
                shares = 0
        
        # 最终价值
        final_price = stock_data[-1]['close']
        final_value = capital + shares * final_price
        profit = final_value - initial_capital
        profit_rate = (profit / initial_capital) * 100
        
        # 买入持有
        buy_hold_shares = int(initial_capital / stock_data[0]['close'])
        buy_hold_value = buy_hold_shares * final_price
        buy_hold_rate = ((buy_hold_value - initial_capital) / initial_capital) * 100
        
        # 胜率
        total_closed_trades = win_trades + loss_trades
        win_rate = (win_trades / total_closed_trades * 100) if total_closed_trades > 0 else 0
        
        return {
            'initial_capital': initial_capital,
            'final_value': final_value,
            'profit': profit,
            'profit_rate': profit_rate,
            'trades_count': len(trades),
            'win_trades': win_trades,
            'loss_trades': loss_trades,
            'win_rate': win_rate,
            'avg_profit_per_trade': (total_profit / total_closed_trades * 100) if total_closed_trades > 0 else 0,
            'trades': trades,
            'buy_hold_rate': buy_hold_rate,
            'excess_return': profit_rate - buy_hold_rate
        }

def main():
    """测试"""
    print("="*90)
    print("📊 带止盈策略的高胜率预测算法")
    print("="*90)
    print("\n策略说明:")
    print("  ✅ 买入条件: 多指标综合评分≥8分，信心度≥70%")
    print("  ✅ 卖出条件: 盈利达到8%自动止盈，或从最高点回撤3%止盈")
    print("  ✅ 同一股票: 买入 -> 持有 -> 止盈卖出 -> 寻找下次买入")
    print("  ✅ 胜率统计: 只要卖出时盈利就算胜")
    print("="*90)
    
    # 生成测试数据（模拟有波动的行情）
    np.random.seed(42)
    dates = pd.date_range(start='2024-01-01', periods=150, freq='D')
    
    # 生成有趋势的价格数据
    base_price = 100
    prices = [base_price]
    
    for i in range(149):
        # 添加趋势和波动
        trend = 0.05 if i < 50 else (-0.05 if i < 100 else 0.03)
        noise = np.random.randn() * 2
        change = trend + noise
        new_price = prices[-1] + change
        prices.append(max(new_price, 50))
    
    stock_data = []
    for i, date in enumerate(dates):
        stock_data.append({
            'date': date.strftime('%Y-%m-%d'),
            'open': prices[i] * (1 + np.random.uniform(-0.01, 0.01)),
            'close': prices[i],
            'high': prices[i] * (1 + abs(np.random.uniform(0, 0.02))),
            'low': prices[i] * (1 - abs(np.random.uniform(0, 0.02))),
            'volume': np.random.randint(10000, 100000)
        })
    
    # 创建预测器
    predictor = ProfitTargetPredictor(
        min_buy_score=8,
        min_confidence=70,
        profit_target=0.08,  # 8%止盈
        trailing_stop=0.03   # 3%回撤止盈
    )
    
    # 分析
    print("\n📈 分析交易信号...")
    trading_points = predictor.analyze(stock_data)
    
    print(f"\n找到 {len(trading_points)} 个交易信号:\n")
    if trading_points:
        print(f"{'日期':<12} {'操作':<12} {'价格':<10} {'信心度':<8} {'盈利率':<10} {'原因':<50}")
        print("-" * 105)
        
        for point in trading_points:
            profit_str = f"{point.profit_rate*100:+.2f}%" if point.profit_rate else "-"
            reasons_str = " | ".join(point.reasons)
            print(f"{point.date:<12} {point.signal.value:<12} {point.price:<10.2f} "
                  f"{point.confidence:<8.1f} {profit_str:<10} {reasons_str:<50}")
    
    # 当前建议
    print("\n" + "="*90)
    print("💡 当前交易建议")
    print("="*90)
    
    recommendation = predictor.get_current_recommendation(stock_data)
    print(f"\n信号: {recommendation['signal_cn']}")
    print(f"信心度: {recommendation['confidence']:.1f}%")
    print(f"价格: ¥{recommendation['price']:.2f}")
    if recommendation.get('profit_rate'):
        print(f"当前盈利: {recommendation['profit_rate']:.2f}%")
    print(f"原因:")
    for reason in recommendation['reasons']:
        print(f"  • {reason}")
    
    # 回测
    print("\n" + "="*90)
    print("🔬 策略回测")
    print("="*90)
    
    backtest = predictor.backtest(stock_data, 100000)
    
    print(f"\n📊 回测结果:")
    print(f"  初始资金: ¥{backtest['initial_capital']:,.0f}")
    print(f"  最终价值: ¥{backtest['final_value']:,.0f}")
    print(f"  总收益: ¥{backtest['profit']:,.0f}")
    print(f"  收益率: {backtest['profit_rate']:.2f}%")
    print(f"  交易次数: {backtest['trades_count']}")
    print(f"  盈利交易: {backtest['win_trades']}")
    print(f"  亏损交易: {backtest['loss_trades']}")
    print(f"  胜率: {backtest['win_rate']:.1f}%")
    print(f"  平均每笔收益: {backtest['avg_profit_per_trade']:.2f}%")
    
    print(f"\n📈 对比基准:")
    print(f"  买入持有收益率: {backtest['buy_hold_rate']:.2f}%")
    print(f"  超额收益: {backtest['excess_return']:.2f}%")
    
    if backtest['trades']:
        print(f"\n📝 交易记录:")
        print(f"{'日期':<12} {'操作':<6} {'价格':<10} {'数量':<8} {'盈利率':<10} {'原因':<40}")
        print("-" * 95)
        
        for trade in backtest['trades']:
            profit_str = f"{trade.get('profit_rate', 0):+.2f}%" if 'profit_rate' in trade else "-"
            reasons_str = " | ".join(trade['reasons'][:2])
            print(f"{trade['date']:<12} {trade['action']:<6} {trade['price']:<10.2f} "
                  f"{trade['shares']:<8} {profit_str:<10} {reasons_str:<40}")
    
    # 导出
    output = {
        'strategy': 'profit_target',
        'profit_target': predictor.profit_target * 100,
        'trailing_stop': predictor.trailing_stop * 100,
        'trading_points': [
            {
                'date': p.date,
                'price': p.price,
                'signal': p.signal.value,
                'signal_type': p.signal_type.value,
                'confidence': p.confidence,
                'reasons': p.reasons,
                'profit_rate': p.profit_rate * 100 if p.profit_rate else 0
            }
            for p in trading_points
        ],
        'backtest': backtest
    }
    
    with open('profit_target_prediction.json', 'w', encoding='utf-8') as f:
        json.dump(output, f, ensure_ascii=False, indent=2)
    
    print(f"\n💾 结果已保存到 profit_target_prediction.json")
    print("\n" + "="*90)

if __name__ == "__main__":
    main()


